#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@file    dataloader.py
@brief
@details
@author  Shivelino
@date    2023-12-23 19:08
@version 0.0.1

@par Copyright(c):
@par todo:
@par history:
"""
import torch
import torchvision
import torchvision.transforms as transforms


def get_dataloader_mnist(data_dir: str, batch_size: int):
    """get MNIST dataloader"""
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    # 训练集
    trainset = torchvision.datasets.MNIST(root=data_dir, train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
    # 测试集
    testset = torchvision.datasets.MNIST(root=data_dir, train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)
    return trainloader, testloader
